#!/bin/bash
# PostgreSQL备份恢复与修复一体化脚本 (增强版)
# 用法: ./restore-and-repair-postgres.sh

# 设置彩色输出
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m' # 恢复默认颜色

# 从.env文件加载配置
load_env() {
  if [ -f ".env" ]; then
    echo "从当前目录加载.env文件..."
    source .env
  elif [ -f "../.env" ]; then
    echo "从上级目录加载.env文件..."
    source ../.env
  else
    echo "警告: 找不到.env文件，使用默认配置"
  fi
}

load_env

# 配置参数 (优先使用.env中的值)
SERVER_TYPE=${SERVER_TYPE:-"B"}
CONTAINER_NAME="postgres-${SERVER_TYPE}"
PGBACKREST_CONTAINER="pgbackrest-${SERVER_TYPE}"
PG_DATA_DIR="/var/lib/postgresql/data/pgdata"
PG_PARENT_DIR="/var/lib/postgresql/data"
BACKUP_DIR=${PGBACKREST_BACKUP_DIR:-"/backup/pgbackrest"}
STANZA_NAME="poddb"
LOG_LEVEL=${PGBACKREST_LOG_LEVEL:-"info"}
POSTGRES_IMAGE=${POSTGRES_IMAGE:-"registry.cn-beijing.aliyuncs.com/xscha/base:postgres_14_4"}
HOST_DATA_DIR=${POSTGRES_DATA_DIR:-"/home/docker1/postgresql/data"}

# 打印带颜色的标题
print_header() {
  echo -e "\n${BLUE}===== $1 =====${NC}\n"
}

# 打印带颜色的成功消息
print_success() {
  echo -e "${GREEN}✅ $1${NC}"
}

# 打印带颜色的警告消息
print_warning() {
  echo -e "${YELLOW}⚠️ $1${NC}"
}

# 打印带颜色的错误消息
print_error() {
  echo -e "${RED}❌ $1${NC}"
}

# 检查命令执行状态
check_status() {
  if [ $? -eq 0 ]; then
    print_success "$1"
    return 0
  else
    print_error "$2"
    return 1
  fi
}

# 检查是否以root身份运行
if [ "$(id -u)" != "0" ]; then
  print_error "此脚本必须以root身份运行"
  echo "请使用sudo运行此脚本:"
  echo "sudo $0"
  exit 1
fi

print_header "PostgreSQL 备份恢复与修复一体化工具 (增强版)"

echo "当前配置:"
echo "  服务器类型: ${SERVER_TYPE}"
echo "  PostgreSQL容器: ${CONTAINER_NAME}"
echo "  pgBackRest容器: ${PGBACKREST_CONTAINER}"
echo "  备份目录: ${BACKUP_DIR}"
echo "  主机数据目录: ${HOST_DATA_DIR}"
echo "  PostgreSQL镜像: ${POSTGRES_IMAGE}"

# 检查必要的目录和容器
echo "检查环境..."
if ! docker ps -a | grep -q ${PGBACKREST_CONTAINER}; then
  print_error "找不到${PGBACKREST_CONTAINER}容器，请确保您的Docker环境已正确设置"
  exit 1
fi

if ! docker ps -a | grep -q ${CONTAINER_NAME}; then
  print_error "找不到${CONTAINER_NAME}容器，请确保您的Docker环境已正确设置"
  exit 1
fi

if [ ! -d "${BACKUP_DIR}" ]; then
  print_error "备份目录 ${BACKUP_DIR} 不存在，请先创建"
  echo "mkdir -p ${BACKUP_DIR}"
  exit 1
fi

echo "检查备份文件..."
if [ ! -d "${BACKUP_DIR}/backup" ]; then
  print_error "备份目录 ${BACKUP_DIR}/backup 不存在"
  echo "请确保已将备份文件正确解压到 ${BACKUP_DIR}"
  echo "提示: 在服务器B上执行以下命令解压备份:"
  echo "  cd ${BACKUP_DIR}/archive"
  echo "  tar -xzf pgbackup_*.tar.gz -C ${BACKUP_DIR}"
  exit 1
fi

# 检查备份目录是否有效
BACKUP_COUNT=$(find "${BACKUP_DIR}/backup" -maxdepth 1 -type d | wc -l)
if [ "$BACKUP_COUNT" -le 1 ]; then
  print_error "未找到有效的备份目录，请检查备份文件是否正确解压"
  echo "备份目录应包含类似 ${BACKUP_DIR}/backup/20250418-084200F 的子目录"
  exit 1
fi

print_success "找到了 $((BACKUP_COUNT-1)) 个备份目录"

# 确认主机数据目录
if [ ! -d "${HOST_DATA_DIR}" ]; then
  print_warning "主机数据目录 ${HOST_DATA_DIR} 不存在，将创建此目录"
  mkdir -p "${HOST_DATA_DIR}"
fi

# 询问用户是否继续
echo ""
echo "此脚本将执行以下操作:"
echo "1. 停止PostgreSQL容器"
echo "2. 完全重置数据目录"
echo "3. 执行备份恢复操作"
echo "4. 重新创建并启动PostgreSQL容器"
echo ""
echo "确保您已备份重要数据，此过程将覆盖现有的PostgreSQL数据"
read -p "是否继续? (y/n): " -n 1 -r
echo ""
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
  echo "操作已取消"
  exit 0
fi

# 第1步: 停止并删除PostgreSQL容器
print_header "1. 停止并备份PostgreSQL容器配置"
echo "停止 ${CONTAINER_NAME} 容器..."
docker stop ${CONTAINER_NAME}
check_status "PostgreSQL服务已停止" "无法停止PostgreSQL服务，但将继续尝试"
sleep 2

# 保存容器配置以便稍后重建
echo "保存容器配置..."
CONTAINER_IMAGE=$(docker inspect -f '{{.Config.Image}}' ${CONTAINER_NAME} || echo "${POSTGRES_IMAGE}")
print_success "容器镜像: $CONTAINER_IMAGE"

# 获取容器配置
ENV_VARS=$(docker inspect -f '{{range .Config.Env}}--env {{.}} {{end}}' ${CONTAINER_NAME} || echo "--env POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-thingbelt2024}")
VOLUMES=$(docker inspect -f '{{range .Mounts}}--volume {{.Source}}:{{.Destination}}{{if .Mode}}:{{.Mode}}{{end}} {{end}}' ${CONTAINER_NAME} || echo "--volume ${HOST_DATA_DIR}:/var/lib/postgresql/data")
PORTS=$(docker inspect -f '{{range $p, $conf := .NetworkSettings.Ports}}{{if $conf}}--publish {{(index $conf 0).HostPort}}:{{$p}} {{end}}{{end}}' ${CONTAINER_NAME} || echo "--publish 5432:5432")
NETWORKS=$(docker inspect -f '{{range $net, $conf := .NetworkSettings.Networks}}--network {{$net}} {{end}}' ${CONTAINER_NAME} || echo "")

# 备份现有容器
echo "重命名现有容器为 ${CONTAINER_NAME}_old..."
docker rename ${CONTAINER_NAME} ${CONTAINER_NAME}_old 2>/dev/null || true

# 第2步: 确保pgbackrest容器运行
print_header "2. 确保pgBackRest容器运行"
if ! docker ps | grep -q ${PGBACKREST_CONTAINER}; then
  echo "启动 ${PGBACKREST_CONTAINER} 容器..."
  docker start ${PGBACKREST_CONTAINER}
  check_status "${PGBACKREST_CONTAINER} 容器已启动" "无法启动 ${PGBACKREST_CONTAINER} 容器"
else
  print_success "${PGBACKREST_CONTAINER} 容器已在运行"
fi

# 第3步: 完全重置数据目录
print_header "3. 完全重置数据目录"
echo "将完全清除并重建数据目录: ${HOST_DATA_DIR}"
echo "此操作将删除所有现有数据!"
read -p "确定要清除数据目录吗? (y/n): " -n 1 -r
echo ""
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
  print_warning "跳过数据目录重置，继续下一步"
else
  echo "重置数据目录..."

  # 创建备份
  BACKUP_TIME=$(date +%Y%m%d%H%M%S)
  if [ -d "${HOST_DATA_DIR}" ] && [ "$(ls -A ${HOST_DATA_DIR} 2>/dev/null)" ]; then
    print_warning "创建现有数据的备份..."
    mkdir -p "/tmp/pg_data_backup_${BACKUP_TIME}"
    cp -a "${HOST_DATA_DIR}/." "/tmp/pg_data_backup_${BACKUP_TIME}/" 2>/dev/null || true
    print_success "数据已备份到 /tmp/pg_data_backup_${BACKUP_TIME}"
  fi

  # 清除数据目录
  rm -rf "${HOST_DATA_DIR}"/*
  mkdir -p "${HOST_DATA_DIR}/pgdata"
  mkdir -p "${HOST_DATA_DIR}/pgdata/pg_wal"
  mkdir -p "${HOST_DATA_DIR}/pgdata/log"

  # 设置权限 (确保使用固定的postgres用户ID 999和组ID 999)
  chown -R 999:999 "${HOST_DATA_DIR}"
  chmod -R 700 "${HOST_DATA_DIR}"

  print_success "数据目录已重置"
fi

# 第4步: 执行备份恢复
print_header "4. 执行备份恢复操作"

echo "准备pgBackRest配置..."
# 在pgbackrest容器中创建配置目录
docker exec -u root ${PGBACKREST_CONTAINER} bash -c "mkdir -p /etc/pgbackrest && chmod 750 /etc/pgbackrest"

# 创建最小化pgbackrest.conf配置
docker exec -u root ${PGBACKREST_CONTAINER} bash -c "cat > /etc/pgbackrest/pgbackrest.conf << EOF
[global]
log-path=/var/log/pgbackrest
repo1-path=${BACKUP_DIR}
process-max=4
delta=y
compress-level=6

[${STANZA_NAME}]
pg1-path=${PG_DATA_DIR}
EOF"

# 确保日志目录存在
docker exec -u root ${PGBACKREST_CONTAINER} bash -c "mkdir -p /var/log/pgbackrest && chown -R postgres:postgres /var/log/pgbackrest && chmod 750 /var/log/pgbackrest"

# 创建stanza
echo "创建pgBackRest存储桶..."
docker exec ${PGBACKREST_CONTAINER} pgbackrest --stanza=${STANZA_NAME} stanza-create || true

# 执行恢复操作
echo "开始恢复操作，此过程可能需要几分钟..."
docker exec ${PGBACKREST_CONTAINER} pgbackrest --stanza=${STANZA_NAME} --log-level-console=${LOG_LEVEL} --pg1-path=${PG_DATA_DIR} --repo1-path=${BACKUP_DIR} --delta restore

if [ $? -ne 0 ]; then
  print_warning "标准恢复失败，尝试使用--force参数"
  docker exec ${PGBACKREST_CONTAINER} pgbackrest --stanza=${STANZA_NAME} --log-level-console=${LOG_LEVEL} --pg1-path=${PG_DATA_DIR} --repo1-path=${BACKUP_DIR} --delta --force restore

  if [ $? -ne 0 ]; then
    print_error "恢复操作失败，无法继续"
    echo "请检查pgBackRest日志以获取更多信息:"
    docker exec ${PGBACKREST_CONTAINER} cat /var/log/pgbackrest/pgbackrest.log 2>/dev/null | tail -n 50 || echo "无法读取日志"
    exit 1
  fi
fi

print_success "备份恢复操作已完成"

# 第5步: 优化PostgreSQL配置
print_header "5. 优化PostgreSQL配置"

# 确保日志目录存在并有正确权限
docker exec -u root ${PGBACKREST_CONTAINER} bash -c "mkdir -p ${PG_DATA_DIR}/log && chown postgres:postgres ${PG_DATA_DIR}/log && chmod 700 ${PG_DATA_DIR}/log"

# 修改配置文件
docker exec -u root ${PGBACKREST_CONTAINER} bash -c "cat > ${PG_DATA_DIR}/postgresql.auto.conf << EOF
# 自动优化配置 - $(date)
listen_addresses = '*'
hot_standby = on
wal_level = replica
archive_mode = off
logging_collector = on
log_directory = 'log'
log_filename = 'postgresql-%Y-%m-%d_%H%M%S.log'
log_statement = 'all'
log_min_error_statement = error
log_min_messages = warning
log_line_prefix = '%m [%p] '
EOF"

# 修改pg_hba.conf允许所有连接
docker exec -u root ${PGBACKREST_CONTAINER} bash -c "cat > ${PG_DATA_DIR}/pg_hba.conf << EOF
# 允许所有连接 (添加于 $(date))
local   all             all                                     trust
host    all             all             127.0.0.1/32            trust
host    all             all             ::1/128                 trust
host    all             all             0.0.0.0/0               md5
EOF"

# 修复权限
docker exec -u root ${PGBACKREST_CONTAINER} bash -c "chown -R postgres:postgres ${PG_DATA_DIR} && chmod 700 ${PG_DATA_DIR} && chmod 600 ${PG_DATA_DIR}/postgresql.auto.conf ${PG_DATA_DIR}/pg_hba.conf"

# 删除恢复信号文件
docker exec -u root ${PGBACKREST_CONTAINER} bash -c "rm -f ${PG_DATA_DIR}/recovery.signal ${PG_DATA_DIR}/standby.signal ${PG_DATA_DIR}/postmaster.pid"

print_success "PostgreSQL配置已优化"

# 第6步: 在宿主机上修复权限
print_header "6. 在宿主机上修复数据目录权限"

echo "确保主机上的数据目录权限正确..."
# 重新设置主机目录的权限
chown -R 999:999 "${HOST_DATA_DIR}"
chmod -R 700 "${HOST_DATA_DIR}"

if [ -d "${HOST_DATA_DIR}/pgdata" ]; then
  find "${HOST_DATA_DIR}/pgdata" -type f -exec chmod 600 {} \;
  find "${HOST_DATA_DIR}/pgdata" -type d -exec chmod 700 {} \;
fi

print_success "主机数据目录权限已修复"

# 第7步: 创建新的PostgreSQL容器
print_header "7. 创建新的PostgreSQL容器"
echo "使用以下配置创建新容器:"
echo "镜像: $CONTAINER_IMAGE"
echo "环境变量: $ENV_VARS"
echo "卷: $VOLUMES"
echo "端口: $PORTS"
echo "网络: $NETWORKS"

# 构建创建命令
CREATE_CMD="docker run -d --name ${CONTAINER_NAME} $ENV_VARS $VOLUMES $PORTS $NETWORKS $CONTAINER_IMAGE"
echo "执行: $CREATE_CMD"
eval $CREATE_CMD

if [ $? -ne 0 ]; then
  print_error "创建新容器失败"
  echo "尝试使用简化配置创建容器..."

  # 使用简化配置
  docker run -d --name ${CONTAINER_NAME} \
    -e POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-thingbelt2024} \
    -v ${HOST_DATA_DIR}:/var/lib/postgresql/data \
    -p 5432:5432 \
    ${CONTAINER_IMAGE}

  if [ $? -ne 0 ]; then
    print_error "无法创建PostgreSQL容器，恢复原始容器"
    docker rename ${CONTAINER_NAME}_old ${CONTAINER_NAME} 2>/dev/null || true
    docker start ${CONTAINER_NAME}
    exit 1
  fi
fi

print_success "新容器已创建"
echo "旧容器 (${CONTAINER_NAME}_old) 可安全删除"

# 第8步: 验证PostgreSQL是否正常工作
print_header "8. 验证PostgreSQL状态"
echo "等待PostgreSQL启动 (最多60秒)..."

for i in {1..30}; do
  if docker ps | grep -q "${CONTAINER_NAME}.*healthy"; then
    print_success "PostgreSQL已成功启动!"
    break
  elif docker ps | grep -q "${CONTAINER_NAME}.*unhealthy\|${CONTAINER_NAME}.*Restarting"; then
    print_warning "PostgreSQL启动出现问题，查看日志..."
    docker logs ${CONTAINER_NAME} | tail -n 30
    break
  else
    echo "等待启动... ($i/30)"
    sleep 2
  fi
done

echo "检查PostgreSQL连接状态..."
sleep 5 # 给PostgreSQL一点额外时间启动

if docker exec ${CONTAINER_NAME} pg_isready -h localhost -p 5432 -U postgres >/dev/null 2>&1; then
  print_success "PostgreSQL服务正常运行"
  echo "尝试执行一个简单查询:"
  docker exec ${CONTAINER_NAME} psql -U postgres -c "SELECT version(), current_timestamp;" || true
else
  print_warning "PostgreSQL服务可能未正常运行"
  echo "查看服务状态和最新日志:"
  docker ps | grep ${CONTAINER_NAME}
  docker logs ${CONTAINER_NAME} | tail -n 30

  print_header "9. 尝试最后的修复措施"
  # 停止容器
  docker stop ${CONTAINER_NAME}

  # 重置数据目录所有权
  echo "完全重置数据目录所有权..."
  chown -R 0:0 "${HOST_DATA_DIR}"  # 先设为root
  chmod -R 700 "${HOST_DATA_DIR}"
  chown -R 999:999 "${HOST_DATA_DIR}"  # 再设为postgres
  chmod -R 700 "${HOST_DATA_DIR}"

  # 启动容器并观察
  docker start ${CONTAINER_NAME}
  echo "等待最终启动 (30秒)..."

  for i in {1..15}; do
    if docker exec ${CONTAINER_NAME} pg_isready -h localhost -p 5432 -U postgres >/dev/null 2>&1; then
      print_success "PostgreSQL在最终修复后成功启动!"
      docker exec ${CONTAINER_NAME} psql -U postgres -c "SELECT version(), current_timestamp;" || true
      break
    else
      echo "等待启动... ($i/15)"
      sleep 2
    fi
  done

  if ! docker exec ${CONTAINER_NAME} pg_isready -h localhost -p 5432 -U postgres >/dev/null 2>&1; then
    print_error "所有修复尝试均失败"
    echo "请考虑以下手动修复步骤:"
    echo "1. 检查容器卷映射是否正确:"
    docker inspect ${CONTAINER_NAME} | grep -A 10 Mounts
    echo "2. 检查数据目录是否有效:"
    ls -la "${HOST_DATA_DIR}/pgdata" 2>/dev/null || echo "无法访问 ${HOST_DATA_DIR}/pgdata"
    echo "3. 以调试模式启动PostgreSQL:"
    echo "   docker exec -it -u postgres ${PGBACKREST_CONTAINER} /usr/lib/postgresql/14/bin/postgres -D ${PG_DATA_DIR} -d 5"
  fi
fi

print_header "PostgreSQL 备份恢复与修复操作完成"
echo "备份恢复和修复操作已完成。"
echo ""
echo "如果成功启动，请执行以下操作以完成配置:"
echo "1. 设置归档模式 (如需要):"
echo "   docker exec -it ${CONTAINER_NAME} psql -U postgres -c \"ALTER SYSTEM SET archive_mode = 'on';\""
echo "   docker exec -it ${CONTAINER_NAME} psql -U postgres -c \"ALTER SYSTEM SET archive_command = 'pgbackrest --stanza=${STANZA_NAME} archive-push %p';\""
echo "   docker exec -it ${CONTAINER_NAME} psql -U postgres -c \"SELECT pg_reload_conf();\""
echo ""
echo "2. 验证数据库内容是否正确:"
echo "   docker exec -it ${CONTAINER_NAME} psql -U postgres -c \"\\l\""
echo ""
echo "3. 如果一切正常，可以删除旧容器:"
echo "   docker rm ${CONTAINER_NAME}_old"
echo ""
echo "4. 如果启动失败，请参考以下调试指南:"
echo "   https://www.postgresql.org/docs/14/startup-failure.html"
echo ""
echo "脚本完成。"
